/*
 * Copyright 2013 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "rva.h"

#include <stddef.h>
#include <string.h>

#include "reimport.h"

BOOL InitRVAPointer(rva_pointer* p, pe_sections* peSections, int rva) {
	p->peSections = peSections;
	p->rva = 0;
	p->data = NULL;
	p->sectStart = 0;
	p->sectEnd = 0;
	return SetRVA(p, rva);
}

BOOL SetRVA(rva_pointer* p, int rva) {
	int sect;
	int start;
	int size;
	for (sect = 0; sect < p->peSections->numberOfSections; sect++) {
		start = p->peSections->sections[sect].VirtualAddress;
		if (rva < start) {
			continue;
		}
		size = p->peSections->sections[sect].Misc.VirtualSize;
		if (size == 0) {
			size = p->peSections->sections[sect].SizeOfRawData;
		}
		size = AlignUp(size, p->peSections->sectionAlignment);
		if (rva >= start + size) {
			continue;
		}
		p->sectStart = start;
		p->sectEnd = start + size;
		p->fileEnd = start
				+ AlignUp(p->peSections->sections[sect].SizeOfRawData,
						p->peSections->fileAlignment);
		if (p->fileEnd > p->sectEnd) {
			p->fileEnd = p->sectEnd;
		}
		p->rva = rva;
		p->data = p->peSections->mapping
				+ p->peSections->sections[sect].PointerToRawData
				+ (p->rva - p->sectStart);
		return TRUE;
	}
	return FALSE;
}

DWORD ReadByte(rva_pointer* p) {
	if (p->rva < p->fileEnd) {
		return *(unsigned char*) p->data;
	}
	return 0;
}

BOOL ReadDWORD(rva_pointer* p, DWORD* data) {
	rva_pointer pCopy;
	int i;
	if (p->rva + sizeof(DWORD) <= p->fileEnd) {
		*data = *(DWORD*) p->data;
		return TRUE;
	}
	pCopy = *p;
	*data = 0;
	for (i = 0; i < sizeof(DWORD); i++) {
		if (!SetRVA(&pCopy, p->rva + i)) {
			return FALSE;
		}
		*data += ReadByte(&pCopy) << (8 * i);
	}
	return TRUE;
}

char* ReadString(rva_pointer* p) {
	rva_pointer pCopy;
	int i;
	char* result;
	char* result_copy;
	int cap;
	pCopy = *p;
	i = 0;
	result = malloc(10);
	cap = 10;
	while (1) {
		result[i] = ReadByte(&pCopy);
		if (result[i] == 0) {
			break;
		}
		i++;
		if (i == cap) {
			cap = 2 * cap;
			result_copy = realloc(result, cap);
			if (result_copy == NULL) {
				free(result);
				return NULL;
			}
			result = result_copy;
		}
		if (!SetRVA(&pCopy, pCopy.rva + 1)) {
			free(result);
			return NULL;
		}
	}
	return result;
}

BOOL ReadData(rva_pointer* p, void* data, int size) {
	rva_pointer pCopy;
	unsigned char* dataCopy;
	int i;
	dataCopy = (unsigned char*) data;
	if (p->rva + size <= p->fileEnd) {
		memcpy(data, p->data, size);
		return TRUE;
	}
	pCopy = *p;
	for (i = 0; i < size; i++) {
		if (!SetRVA(&pCopy, p->rva + i)) {
			return FALSE;
		}
		*(unsigned char*) dataCopy = ReadByte(&pCopy);
		dataCopy++;
	}
	return TRUE;
}
